#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

红队生态联动接口
支持API集成和C2平台插件体系
"""

import json
import time
import base64
import hashlib
import hmac
import requests
import threading
import logging
from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass, asdict
from enum import Enum
from abc import ABC, abstractmethod
import socket
import ssl
from urllib.parse import urljoin
import subprocess
import os




class IntegrationType(Enum):
    """集成类型枚举"""
    API = "api"
    WEBHOOK = "webhook"
    PLUGIN = "plugin"
    SCRIPT = "script"
    C2_BEACON = "c2_beacon"

class OperationStatus(Enum):
    """操作状态枚举"""
    SUCCESS = "success"
    FAILED = "failed"
    PENDING = "pending"
    TIMEOUT = "timeout"
    ERROR = "error"

@dataclass
class IntegrationConfig:
    """集成配置"""
    name: str
    type: IntegrationType
    endpoint: str
    auth_token: Optional[str] = None
    api_key: Optional[str] = None
    secret_key: Optional[str] = None
    timeout: int = 30
    retry_count: int = 3
    enabled: bool = True
    custom_headers: Dict[str, str] = None
    
    def __post_init__(self):
        if self.custom_headers is None:
            self.custom_headers = {}

@dataclass
class OperationEvent:
    """操作事件"""
    event_id: str
    event_type: str
    timestamp: float
    source: str
    target: str
    operation: str
    status: OperationStatus
    data: Dict[str, Any]
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}

class BaseIntegration(ABC):
    """基础集成接口"""
    
    def __init__(self, config: IntegrationConfig, logger: Optional[logging.Logger] = None):
        self.config = config
        self.logger = logger or logging.getLogger(__name__)
        self.is_connected = False
        
    @abstractmethod
    def connect(self) -> bool:
        """建立连接"""
        pass
    
    @abstractmethod
    def disconnect(self) -> None:
        """断开连接"""
        pass
    
    @abstractmethod
    def send_event(self, event: OperationEvent) -> bool:
        """发送事件"""
        pass
    
    @abstractmethod
    def receive_command(self) -> Optional[Dict[str, Any]]:
        """接收命令"""
        pass

class APIIntegration(BaseIntegration):
    """API集成"""
    
    def __init__(self, config: IntegrationConfig, logger: Optional[logging.Logger] = None):
        super().__init__(config, logger)
        self.session = requests.Session()
        self._setup_session()
    
    def _setup_session(self):
        """设置会话"""
        headers = {
            'Content-Type': 'application/json',
            'User-Agent': 'RedTeam-Integration/1.0'
        }
        
        if self.config.auth_token:
            headers['Authorization'] = "Bearer {}".format(self.config.auth_token)
        elif self.config.api_key:
            headers['X-API-Key'] = self.config.api_key
        
        headers.update(self.config.custom_headers)
        self.session.headers.update(headers)
    
    def connect(self) -> bool:
        """建立连接"""
        try:
            # 测试连接
            response = self.session.get(
                urljoin(self.config.endpoint, '/health'),
                timeout=self.config.timeout
            )
            
            self.is_connected = response.status_code == 200
            if self.is_connected:
                self.logger.info("API集成 {} 连接成功".format(self.config.name))
            else:
                self.logger.error("API集成 {} 连接失败: {}".format(self.config.name, response.status_code))
            
            return self.is_connected
            
        except Exception as e:
            self.logger.error("API集成 {} 连接异常: {}".format(self.config.name, e))
            return False
    
    def disconnect(self) -> None:
        """断开连接"""
        self.session.close()
        self.is_connected = False
        self.logger.info("API集成 {} 已断开".format(self.config.name))
    
    def send_event(self, event: OperationEvent) -> bool:
        """发送事件"""
        try:
            data = asdict(event)
            
            # 添加签名
            if self.config.secret_key:
                data['signature'] = self._generate_signature(data)
            
            response = self.session.post(
                urljoin(self.config.endpoint, '/events'),
                json=data,
                timeout=self.config.timeout
            )
            
            success = response.status_code in [200, 201, 202]
            if success:
                self.logger.debug("事件发送成功: {}".format(event.event_id))
            else:
                self.logger.error("事件发送失败: {} - {}".format(response.status_code, response.text))
            
            return success
            
        except Exception as e:
            self.logger.error("发送事件异常: {}".format(e))
            return False
    
    def receive_command(self) -> Optional[Dict[str, Any]]:
        """接收命令"""
        try:
            response = self.session.get(
                urljoin(self.config.endpoint, '/commands'),
                timeout=self.config.timeout
            )
            
            if response.status_code == 200:
                return response.json()
            elif response.status_code == 204:
                return None  # 无命令
            else:
                self.logger.error("接收命令失败: {}".format(response.status_code))
                return None
                
        except Exception as e:
            self.logger.error("接收命令异常: {}".format(e))
            return None
    
    def _generate_signature(self, data: Dict[str, Any]) -> str:
        """生成签名"""
        message = json.dumps(data, sort_keys=True)
        signature = hmac.new(
            self.config.secret_key.encode(),
            message.encode(),
            hashlib.sha256
        ).hexdigest()
        return signature

class WebhookIntegration(BaseIntegration):
    """Webhook集成"""
    
    def __init__(self, config: IntegrationConfig, logger: Optional[logging.Logger] = None):
        super().__init__(config, logger)
        self.session = requests.Session()
        self._setup_session()
    
    def _setup_session(self):
        """设置会话"""
        headers = {
            'Content-Type': 'application/json',
            'User-Agent': 'RedTeam-Webhook/1.0'
        }
        headers.update(self.config.custom_headers)
        self.session.headers.update(headers)
    
    def connect(self) -> bool:
        """建立连接"""
        self.is_connected = True
        self.logger.info("Webhook集成 {} 已准备".format(self.config.name))
        return True
    
    def disconnect(self) -> None:
        """断开连接"""
        self.session.close()
        self.is_connected = False
    
    def send_event(self, event: OperationEvent) -> bool:
        """发送事件"""
        try:
            payload = {
                'event': asdict(event),
                'timestamp': time.time(),
                'source': 'redteam-system'
            }
            
            # 添加签名
            if self.config.secret_key:
                payload['signature'] = self._generate_signature(payload)
            
            response = self.session.post(
                self.config.endpoint,
                json=payload,
                timeout=self.config.timeout
            )
            
            success = response.status_code in [200, 201, 202]
            if success:
                self.logger.debug("Webhook事件发送成功: {}".format(event.event_id))
            else:
                self.logger.error("Webhook事件发送失败: {}".format(response.status_code))
            
            return success
            
        except Exception as e:
            self.logger.error("Webhook发送异常: {}".format(e))
            return False
    
    def receive_command(self) -> Optional[Dict[str, Any]]:
        """接收命令（Webhook通常是单向的）"""
        return None
    
    def _generate_signature(self, payload: Dict[str, Any]) -> str:
        """生成签名"""
        message = json.dumps(payload, sort_keys=True)
        signature = hmac.new(
            self.config.secret_key.encode(),
            message.encode(),
            hashlib.sha256
        ).hexdigest()
        return "sha256={}".format(signature)

class C2BeaconIntegration(BaseIntegration):
    """C2 Beacon集成"""
    
    def __init__(self, config: IntegrationConfig, logger: Optional[logging.Logger] = None):
        super().__init__(config, logger)
        self.beacon_thread = None
        self.running = False
        self.command_queue = []
        
    def connect(self) -> bool:
        """建立连接"""
        try:
            self.running = True
            self.beacon_thread = threading.Thread(target=self._beacon_loop)
            self.beacon_thread.daemon = True
            self.beacon_thread.start()
            
            self.is_connected = True
            self.logger.info("C2 Beacon {} 已启动".format(self.config.name))
            return True
            
        except Exception as e:
            self.logger.error("C2 Beacon启动失败: {}".format(e))
            return False
    
    def disconnect(self) -> None:
        """断开连接"""
        self.running = False
        if self.beacon_thread:
            self.beacon_thread.join(timeout=5)
        self.is_connected = False
        self.logger.info("C2 Beacon {} 已停止".format(self.config.name))
    
    def send_event(self, event: OperationEvent) -> bool:
        """发送事件"""
        try:
            # 将事件转换为C2格式
            beacon_data = self._format_for_c2(event)
            
            # 发送到C2服务器
            response = requests.post(
                self.config.endpoint,
                data=beacon_data,
                headers={'Content-Type': 'application/octet-stream'},
                timeout=self.config.timeout
            )
            
            success = response.status_code == 200
            if success:
                self.logger.debug("C2事件发送成功: {}".format(event.event_id))
                
                # 解析响应中的命令
                if response.content:
                    commands = self._parse_c2_response(response.content)
                    self.command_queue.extend(commands)
            
            return success
            
        except Exception as e:
            self.logger.error("C2事件发送异常: {}".format(e))
            return False
    
    def receive_command(self) -> Optional[Dict[str, Any]]:
        """接收命令"""
        if self.command_queue:
            return self.command_queue.pop(0)
        return None
    
    def _beacon_loop(self):
        """Beacon循环"""
        while self.running:
            try:
                # 发送心跳
                heartbeat_event = OperationEvent(
                    event_id="heartbeat_{}".format(int(time.time())),
                    event_type="heartbeat",
                    timestamp=time.time(),
                    source="redteam-system",
                    target="c2-server",
                    operation="heartbeat",
                    status=OperationStatus.SUCCESS,
                    data={"status": "alive"}
                )
                
                self.send_event(heartbeat_event)
                
                # 等待下次心跳
                time.sleep(60)  # 1分钟心跳间隔
                
            except Exception as e:
                self.logger.error("Beacon循环异常: {}".format(e))
                time.sleep(30)
    
    def _format_for_c2(self, event: OperationEvent) -> bytes:
        """格式化为C2格式"""
        # 简化的C2格式，实际实现需要根据具体C2协议
        data = json.dumps(asdict(event))
        encoded = base64.b64encode(data.encode())
        return encoded
    
    def _parse_c2_response(self, response: bytes) -> List[Dict[str, Any]]:
        """解析C2响应"""
        try:
            decoded = base64.b64decode(response)
            data = json.loads(decoded.decode())
            
            if isinstance(data, list):
                return data
            elif isinstance(data, dict):
                return [data]
            else:
                return []
                
        except Exception as e:
            self.logger.error("解析C2响应失败: {}".format(e))
            return []

class ScriptIntegration(BaseIntegration):
    """脚本集成"""
    
    def __init__(self, config: IntegrationConfig, logger: Optional[logging.Logger] = None):
        super().__init__(config, logger)
        self.script_path = config.endpoint
    
    def connect(self) -> bool:
        """建立连接"""
        if os.path.exists(self.script_path) and os.access(self.script_path, os.X_OK):
            self.is_connected = True
            self.logger.info("脚本集成 {} 已准备".format(self.config.name))
            return True
        else:
            self.logger.error("脚本文件不存在或不可执行: {}".format(self.script_path))
            return False
    
    def disconnect(self) -> None:
        """断开连接"""
        self.is_connected = False
    
    def send_event(self, event: OperationEvent) -> bool:
        """发送事件"""
        try:
            # 将事件作为JSON传递给脚本
            event_json = json.dumps(asdict(event))
            
            result = subprocess.run(
                [self.script_path, 'event'],
                input=event_json,
                text=True,
                capture_output=True,
                timeout=self.config.timeout
            )
            
            success = result.returncode == 0
            if success:
                self.logger.debug("脚本事件处理成功: {}".format(event.event_id))
            else:
                self.logger.error("脚本事件处理失败: {}".format(result.stderr))
            
            return success
            
        except Exception as e:
            self.logger.error("脚本执行异常: {}".format(e))
            return False
    
    def receive_command(self) -> Optional[Dict[str, Any]]:
        """接收命令"""
        try:
            result = subprocess.run(
                [self.script_path, 'command'],
                capture_output=True,
                text=True,
                timeout=self.config.timeout
            )
            
            if result.returncode == 0 and result.stdout.strip():
                return json.loads(result.stdout.strip())
            
            return None
            
        except Exception as e:
            self.logger.error("脚本命令接收异常: {}".format(e))
            return None

class EcosystemIntegrationManager:
    """生态集成管理器"""
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        self.logger = logger or logging.getLogger(__name__)
        self.integrations: Dict[str, BaseIntegration] = {}
        self.event_handlers: Dict[str, List[Callable]] = {}
        self.running = False
        self.command_thread = None
    
    def add_integration(self, integration: BaseIntegration) -> None:
        """添加集成"""
        self.integrations[integration.config.name] = integration
        self.logger.info("已添加集成: {}".format(integration.config.name))
    
    def remove_integration(self, name: str) -> None:
        """移除集成"""
        if name in self.integrations:
            self.integrations[name].disconnect()
            del self.integrations[name]
            self.logger.info("已移除集成: {}".format(name))
    
    def list_integrations(self) -> List[str]:
        """列出所有集成名称"""
        return list(self.integrations.keys())
    
    def connect_all(self) -> Dict[str, bool]:
        """连接所有集成"""
        results = {}
        
        for name, integration in self.integrations.items():
            if integration.config.enabled:
                try:
                    results[name] = integration.connect()
                except Exception as e:
                    self.logger.error("连接集成 {} 失败: {}".format(name, e))
                    results[name] = False
            else:
                results[name] = False
        
        return results
    
    def disconnect_all(self) -> None:
        """断开所有集成"""
        for integration in self.integrations.values():
            try:
                integration.disconnect()
            except Exception as e:
                self.logger.error("断开集成失败: {}".format(e))
    
    def broadcast_event(self, event: OperationEvent) -> Dict[str, bool]:
        """广播事件到所有集成"""
        results = {}
        
        for name, integration in self.integrations.items():
            if integration.is_connected:
                try:
                    results[name] = integration.send_event(event)
                except Exception as e:
                    self.logger.error("向集成 {} 发送事件失败: {}".format(name, e))
                    results[name] = False
            else:
                results[name] = False
        
        return results
    
    def send_event_to(self, integration_name: str, event: OperationEvent) -> bool:
        """发送事件到指定集成"""
        if integration_name in self.integrations:
            integration = self.integrations[integration_name]
            if integration.is_connected:
                return integration.send_event(event)
        
        return False
    
    def register_event_handler(self, event_type: str, handler: Callable) -> None:
        """注册事件处理器"""
        if event_type not in self.event_handlers:
            self.event_handlers[event_type] = []
        
        self.event_handlers[event_type].append(handler)
        self.logger.info("已注册事件处理器: {}".format(event_type))
    
    def handle_event(self, event: Dict[str, Any]) -> None:
        """处理事件"""
        event_type = event.get('type', 'unknown')
        
        # 调用注册的事件处理器
        if event_type in self.event_handlers:
            for handler in self.event_handlers[event_type]:
                try:
                    handler(event)
                except Exception as e:
                    self.logger.error("事件处理器执行失败: {}".format(e))
        
        # 如果是OperationEvent对象，广播到所有集成
        if hasattr(event, 'event_type'):
            self.broadcast_event(event)
        else:
            # 创建OperationEvent对象并广播
            operation_event = OperationEvent(
                event_id=event.get('id', 'unknown'),
                event_type=event_type,
                timestamp=event.get('timestamp', time.time()),
                source=event.get('source', 'unknown'),
                target=event.get('target', 'unknown'),
                operation=event.get('operation', 'unknown'),
                status=OperationStatus.SUCCESS,
                data=event.get('data', {})
            )
            self.broadcast_event(operation_event)
        
        self.logger.info("已处理事件: {}".format(event_type))
    
    def start_command_listener(self) -> None:
        """启动命令监听器"""
        if not self.running:
            self.running = True
            self.command_thread = threading.Thread(target=self._command_loop)
            self.command_thread.daemon = True
            self.command_thread.start()
            self.logger.info("命令监听器已启动")
    
    def stop_command_listener(self) -> None:
        """停止命令监听器"""
        self.running = False
        if self.command_thread:
            self.command_thread.join(timeout=5)
        self.logger.info("命令监听器已停止")
    
    def _command_loop(self) -> None:
        """命令监听循环"""
        while self.running:
            try:
                for name, integration in self.integrations.items():
                    if integration.is_connected:
                        command = integration.receive_command()
                        if command:
                            self._handle_command(name, command)
                
                time.sleep(5)  # 5秒检查一次
                
            except Exception as e:
                self.logger.error("命令监听循环异常: {}".format(e))
                time.sleep(10)
    
    def _handle_command(self, source: str, command: Dict[str, Any]) -> None:
        """处理命令"""
        try:
            command_type = command.get('type', 'unknown')
            self.logger.info("收到来自 {} 的命令: {}".format(source, command_type))
            
            # 触发事件处理器
            if command_type in self.event_handlers:
                for handler in self.event_handlers[command_type]:
                    try:
                        handler(source, command)
                    except Exception as e:
                        self.logger.error("命令处理器异常: {}".format(e))
            
            # 创建命令事件
            event = OperationEvent(
                event_id="cmd_{}_{}".format(int(time.time()), hash(str(command)) % 10000),
                event_type="command_received",
                timestamp=time.time(),
                source=source,
                target="redteam-system",
                operation=command_type,
                status=OperationStatus.SUCCESS,
                data=command
            )
            
            # 广播命令事件（除了源集成）
            for name, integration in self.integrations.items():
                if name != source and integration.is_connected:
                    integration.send_event(event)
            
        except Exception as e:
            self.logger.error("处理命令异常: {}".format(e))
    
    def create_integration_from_config(self, config: IntegrationConfig) -> BaseIntegration:
        """从配置创建集成"""
        if config.type == IntegrationType.API:
            return APIIntegration(config, self.logger)
        elif config.type == IntegrationType.WEBHOOK:
            return WebhookIntegration(config, self.logger)
        elif config.type == IntegrationType.C2_BEACON:
            return C2BeaconIntegration(config, self.logger)
        elif config.type == IntegrationType.SCRIPT:
            return ScriptIntegration(config, self.logger)
        else:
            raise ValueError("不支持的集成类型: {}".format(config.type))
    
    def load_integrations_from_file(self, filepath: str) -> None:
        """从文件加载集成配置"""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                configs = json.load(f)
            
            for config_data in configs:
                config = IntegrationConfig(**config_data)
                integration = self.create_integration_from_config(config)
                self.add_integration(integration)
            
            self.logger.info("从 {} 加载了 {} 个集成配置".format(filepath, len(configs)))
            
        except Exception as e:
            self.logger.error("加载集成配置失败: {}".format(e))
    
    def get_status(self) -> Dict[str, Any]:
        """获取状态"""
        status = {
            'total_integrations': len(self.integrations),
            'connected_integrations': sum(1 for i in self.integrations.values() if i.is_connected),
            'command_listener_running': self.running,
            'integrations': {}
        }
        
        for name, integration in self.integrations.items():
            status['integrations'][name] = {
                'type': integration.config.type.value,
                'connected': integration.is_connected,
                'enabled': integration.config.enabled,
                'endpoint': integration.config.endpoint
            }
        
        return status

# Cobalt Strike Aggressor Script 兼容接口
class CobaltStrikePlugin:
    """Cobalt Strike插件接口"""
    
    def __init__(self, manager: EcosystemIntegrationManager):
        self.manager = manager
        self.logger = logging.getLogger(__name__)
    
    def beacon_callback(self, bid: str, data: Dict[str, Any]) -> None:
        """Beacon回调"""
        event = OperationEvent(
            event_id="cs_beacon_{}_{}".format(bid, int(time.time())),
            event_type="beacon_callback",
            timestamp=time.time(),
            source="cobalt_strike",
            target=bid,
            operation="beacon_data",
            status=OperationStatus.SUCCESS,
            data=data,
            metadata={"beacon_id": bid}
        )
        
        self.manager.broadcast_event(event)
    
    def execute_command(self, bid: str, command: str) -> bool:
        """执行命令"""
        event = OperationEvent(
            event_id="cs_cmd_{}_{}".format(bid, int(time.time())),
            event_type="command_execution",
            timestamp=time.time(),
            source="cobalt_strike",
            target=bid,
            operation="execute_command",
            status=OperationStatus.PENDING,
            data={"command": command},
            metadata={"beacon_id": bid}
        )
        
        results = self.manager.broadcast_event(event)
        return any(results.values())
    
    def deploy_backdoor(self, bid: str, backdoor_config: Dict[str, Any]) -> bool:
        """部署后门"""
        event = OperationEvent(
            event_id="cs_deploy_{}_{}".format(bid, int(time.time())),
            event_type="backdoor_deployment",
            timestamp=time.time(),
            source="cobalt_strike",
            target=bid,
            operation="deploy_backdoor",
            status=OperationStatus.PENDING,
            data=backdoor_config,
            metadata={"beacon_id": bid}
        )
        
        results = self.manager.broadcast_event(event)
        return any(results.values())

def main():
    """主函数 - 用于测试"""
    logging.basicConfig(level=logging.INFO)
    
    # 创建管理器
    manager = EcosystemIntegrationManager()
    
    # 创建测试集成
    api_config = IntegrationConfig(
        name="test_api",
        type=IntegrationType.API,
        endpoint="http://localhost:8080/api",
        api_key="test_key"
    )
    
    webhook_config = IntegrationConfig(
        name="test_webhook",
        type=IntegrationType.WEBHOOK,
        endpoint="http://localhost:8080/webhook",
        secret_key="test_secret"
    )
    
    # 添加集成
    api_integration = manager.create_integration_from_config(api_config)
    webhook_integration = manager.create_integration_from_config(webhook_config)
    
    manager.add_integration(api_integration)
    manager.add_integration(webhook_integration)
    
    # 连接集成
    results = manager.connect_all()
    print("连接结果: {}".format(results))
    
    # 发送测试事件
    test_event = OperationEvent(
        event_id="test_001",
        event_type="test",
        timestamp=time.time(),
        source="test_system",
        target="test_target",
        operation="test_operation",
        status=OperationStatus.SUCCESS,
        data={"message": "这是一个测试事件"}
    )
    
    broadcast_results = manager.broadcast_event(test_event)
    print("广播结果: {}".format(broadcast_results))
    
    # 获取状态
    status = manager.get_status()
    print("系统状态: {}".format(json.dumps(status, indent=2, ensure_ascii=False)))
    
    # 清理
    manager.disconnect_all()

if __name__ == "__main__":
    main()